import sys
### flank区间要剔除基因区
file1=open(sys.argv[1],'r')# gene_gff
file2=open(sys.argv[2],'r')# TE_gff
file3=open(sys.argv[3],'r')# total_gene_pos

chr_TE_dict={}
for line in file2:
    if line.startswith('#'):
        continue
    line_list=line.strip().split()
    if '(' in line_list[-1]:
        continue
    if line_list[0] not in chr_TE_dict:
        chr_TE_dict[line_list[0]]={}
    if line_list[-1] not in chr_TE_dict[line_list[0]]:
        chr_TE_dict[line_list[0]][line_list[-1]]=[]
    if int(line_list[1])<int(line_list[2]):
        chr_TE_dict[line_list[0]][line_list[-1]].append([int(line_list[1]),int(line_list[2])])
    else:
        chr_TE_dict[line_list[0]][line_list[-1]].append([int(line_list[2]),int(line_list[1])])
    # print(chr_TE_dict)
    # sys.exit()
    
chr_intege_dict={}
total_length=0
for line in file1:
    line_list=line.strip().split('\t')
    if line_list[0] not in chr_intege_dict.keys():
        chr_intege_dict[line_list[0]]=[]
    if int(line_list[1])<int(line_list[2]):
        chr_intege_dict[line_list[0]].append([int(line_list[1]),int(line_list[2])])
        total_length+=int(line_list[2])-int(line_list[1])
    else:
        chr_intege_dict[line_list[0]].append([int(line_list[2]),int(line_list[1])])
        total_length+=int(line_list[1])-int(line_list[2])
    
file1.close()
file1=open(sys.argv[1],'r')
# print('基因区长度\t'+total_length)      
# length=0
out_dict={'CACTA_TIR_transposon':0,'Copia_LTR_retrotransposon':0,'Gypsy_LTR_retrotransposon':0,'hAT_TIR_transposon':0,'helitron':0,'long_terminal_repeat':0,\
'LTR_retrotransposon':0,'Mutator_TIR_transposon':0,'PIF_Harbinger_TIR_transposon':0,'repeat_region':0,'target_site_duplication':0,'Tc1_Mariner_TIR_transposon':0,'TE':0}

def get_inter_dict(chr_TE_dict,chr_intege_dict,out_dict):
    # out_dict={}
    for key0,value0 in chr_TE_dict.items():
        # print(key)
        try:
            gene_list=sorted(chr_intege_dict[key0],key=lambda x:x[0])
        except KeyError:
            continue
        for key1,value1 in value0.items():
            if key1 not in out_dict:
                continue
                # out_dict[key1]=0
            TE_list=sorted(value1,key=lambda x:x[0])
            index=0
            length=0
            add_length=0
            # while index<=len(TE_list):
            for intro_pos in gene_list:
                for te_pos in TE_list[index:]:
                    # print(intro_pos,te_pos)
                    if te_pos[0]>=intro_pos[0]:
                        if te_pos[0]>intro_pos[1]:
                            break
                        elif te_pos[1]<=intro_pos[1]:
                            length+=te_pos[1]-te_pos[0]+1
                        elif te_pos[1]>=intro_pos[1]:
                            length+=intro_pos[1]-te_pos[0]+1
                    elif te_pos[0]<=intro_pos[0]:
                        if te_pos[1]<=intro_pos[0]:
                            index+=1
                            continue
                        
                        elif te_pos[1]<=intro_pos[1]:
                            length+=te_pos[1]-intro_pos[0]+1
                        elif te_pos[1]>=intro_pos[1]:
                            length+=intro_pos[1]-intro_pos[0]+1
                    # print(length,te_pos,intro_pos)
                    # sys.exit()  
                    
                   
                    
            out_dict[key1]+=length
    return out_dict
out_dict1=get_inter_dict(chr_TE_dict,chr_intege_dict,out_dict)
# print('TE 长度\t'+length)
# print(length)
# print(out_dict1)
file4=open(sys.argv[4]+'.inter.count','w')
file4.write('基因区长度\t'+str(total_length)+'\n')
for key,value in out_dict1.items():
    file4.write(key+'\t'+str(value/total_length)+'\n')
up_chr_intege_dict={}
total_length=0
for line in file1:
    line_list=line.strip().split('\t')
    if line_list[0] not in up_chr_intege_dict.keys():
        up_chr_intege_dict[line_list[0]]=[]
    up_chr_intege_dict[line_list[0]].append([int(line_list[1])-5000,int(line_list[1])])
    up_chr_intege_dict[line_list[0]].append([int(line_list[2]),int(line_list[2])+5000])
    total_length+=4000
##添加所有基因的bed用来去除flank区间和基因body的重叠
total_gene_region_dict={}
for line in file3:
    line_list=line.strip().split('\t')
    if line_list[3] not in total_gene_region_dict.keys():
        total_gene_region_dict[line_list[3]]=[]
    total_gene_region_dict[line_list[3]].append([int(line_list[4].split(':')[0]),int(line_list[-1].split(':')[1])])
out_dict={'CACTA_TIR_transposon':0,'Copia_LTR_retrotransposon':0,'Gypsy_LTR_retrotransposon':0,'hAT_TIR_transposon':0,'helitron':0,'long_terminal_repeat':0,\
'LTR_retrotransposon':0,'Mutator_TIR_transposon':0,'PIF_Harbinger_TIR_transposon':0,'repeat_region':0,'target_site_duplication':0,'Tc1_Mariner_TIR_transposon':0,'TE':0}

def get_uniq_list(listA,listB):
    result = []
    for intervalA in listA:
        # 标记A区间是否保留
        keep_intervalA = True
        # 遍历B中的每个区间
        for intervalB in listB:
            # 检查A和B区间是否有重叠
            if intervalA[0] <= intervalB[1] and intervalA[1] >= intervalB[0]:
                # A区间与B区间存在重叠
                # 完全重叠情况下，跳过A区间
                if intervalA[0] >= intervalB[0] and intervalA[1] <= intervalB[1]:
                    keep_intervalA = False
                    break
                # 部分重叠情况下，拆分A区间并去除重叠部分
                elif intervalA[0] < intervalB[0] and intervalA[1] > intervalB[1]:
                    result.append([intervalA[0], intervalB[0]])
                    result.append([intervalB[1], intervalA[1]])
                    keep_intervalA = False
                
                # 重叠的部分在A区间的起始或结束位置
                elif intervalA[0] >= intervalB[0]:
                    intervalA[0] = intervalB[1]
                else:
                    intervalA[1] = intervalB[0]
        
        # 如果A区间未与任何B区间重叠，则保留A区间
        if keep_intervalA:
            result.append(intervalA)
    return result
total_length=0        
up_chr_intege_dict2={}
for key,value in up_chr_intege_dict.items():
    # up_chr_intege_dict2[key]=[]
    out_list=get_uniq_list(value,total_gene_region_dict[key])
    up_chr_intege_dict2[key]=out_list
    for i in out_list:
        total_length+=(i[1]-i[0])

out_dict2=get_inter_dict(chr_TE_dict,up_chr_intege_dict2,out_dict)
file5=open(sys.argv[4]+'.up.count_5K','w')
file5.write('基因上下游长度\t'+str(total_length)+'\n')
for key,value in out_dict2.items():
    file5.write(key+'\t'+str(value/total_length)+'\n')